from itertools import combinations
from collections import defaultdict


class Apriori:
    def __init__(self, min_support=0.5, min_confidence=0.7):
        self.min_support = min_support
        self.min_confidence = min_confidence
        self.frequent_itemsets = []
        self.association_rules = []

    def _get_frequent_1_itemsets(self, transactions):
        """获取频繁1项集"""
        item_counts = defaultdict(int)
        for transaction in transactions:
            for item in transaction:
                item_counts[frozenset([item])] += 1

        n_transactions = len(transactions)
        return {itemset: count / n_transactions
                for itemset, count in item_counts.items()
                if count / n_transactions >= self.min_support}

    def _has_infrequent_subset(self, candidate, prev_frequent, k):
        """检查候选项集是否有非频繁子集"""
        subsets = combinations(candidate, k - 1)
        for subset in subsets:
            if frozenset(subset) not in prev_frequent:
                return True
        return False

    def _generate_candidates(self, prev_frequent, k):
        """生成候选k项集"""
        candidates = set()

        # 将字典键转换为列表以便使用索引
        prev_itemsets = list(prev_frequent.keys())
        n = len(prev_itemsets)

        # 自连接：合并两个(k-1)项集来生成k项集候选
        for i in range(n):
            for j in range(i + 1, n):
                # 检查前k-2项是否相同
                itemset1 = list(prev_itemsets[i])
                itemset2 = list(prev_itemsets[j])
                itemset1.sort()
                itemset2.sort()

                if itemset1[:k - 2] == itemset2[:k - 2]:
                    candidate = prev_itemsets[i] | prev_itemsets[j]
                    if len(candidate) == k and not self._has_infrequent_subset(candidate, prev_frequent, k):
                        candidates.add(candidate)

        return candidates

    def fit(self, transactions):
        """挖掘频繁项集"""
        # 获取频繁1项集
        frequent_1 = self._get_frequent_1_itemsets(transactions)
        self.frequent_itemsets.append(frequent_1)

        k = 2
        while True:
            if len(self.frequent_itemsets[-1]) == 0:
                break

            candidates = self._generate_candidates(self.frequent_itemsets[-1], k)
            if not candidates:
                break

            # 计算候选项集支持度
            candidate_counts = defaultdict(int)
            for transaction in transactions:
                transaction_set = set(transaction)
                for candidate in candidates:
                    if candidate.issubset(transaction_set):
                        candidate_counts[candidate] += 1

            n_transactions = len(transactions)
            frequent_k = {itemset: count / n_transactions
                          for itemset, count in candidate_counts.items()
                          if count / n_transactions >= self.min_support}

            if not frequent_k:
                break

            self.frequent_itemsets.append(frequent_k)
            k += 1

    def generate_rules(self):
        """生成关联规则"""
        for frequent_dict in self.frequent_itemsets[1:]:  # 跳过1项集
            for itemset, support in frequent_dict.items():
                items = list(itemset)
                if len(items) < 2:
                    continue

                # 生成所有可能的规则
                for i in range(1, len(items)):
                    for antecedent in combinations(items, i):
                        antecedent = frozenset(antecedent)
                        consequent = itemset - antecedent

                        if len(consequent) == 0:
                            continue

                        # 找到前项的支持度
                        antecedent_support = 0
                        for freq_dict in self.frequent_itemsets:
                            if antecedent in freq_dict:
                                antecedent_support = freq_dict[antecedent]
                                break

                        if antecedent_support > 0:
                            confidence = support / antecedent_support
                            if confidence >= self.min_confidence:
                                self.association_rules.append({
                                    'antecedent': antecedent,
                                    'consequent': consequent,
                                    'support': support,
                                    'confidence': confidence
                                })


# 示例使用
if __name__ == "__main__":
    # 示例交易数据
    transactions = [
        ['牛奶', '面包', '黄油'],
        ['啤酒', '面包'],
        ['牛奶', '啤酒', '面包'],
        ['面包', '黄油'],
        ['牛奶', '面包']
    ]

    apriori = Apriori(min_support=0.4, min_confidence=0.6)
    apriori.fit(transactions)
    apriori.generate_rules()

    print("频繁项集:")
    for i, freq_dict in enumerate(apriori.frequent_itemsets):
        print(f"{i + 1}项集:", {tuple(itemset): support for itemset, support in freq_dict.items()})

    print("\n关联规则:")
    for rule in apriori.association_rules:
        print(f"{set(rule['antecedent'])} => {set(rule['consequent'])} "
              f"(支持度: {rule['support']:.2f}, 置信度: {rule['confidence']:.2f})")